﻿#include "../exercise.h"
#include <cstring>
// READ: 类模板 <https://zh.cppreference.com/w/cpp/language/class_template>

template<class T>
struct Tensor4D {
    unsigned int shape[4];
    T *data;

    Tensor4D(unsigned int const shape_[4], T const *data_) {
        unsigned int size = 1;
        // TODO: 填入正确的 shape 并计算 size
        for (int i = 0; i < 4; ++i) {
            shape[i] = shape_[i];
            size *= shape[i];
        }

        data = new T[size];
        std::memcpy(data, data_, size * sizeof(T));
    }
    ~Tensor4D() {
        delete[] data;
    }

    // 为了保持简单，禁止复制和移动
    Tensor4D(Tensor4D const &) = delete;
    Tensor4D(Tensor4D &&) noexcept = delete;


    // 计算线性索引
    unsigned int linear_index(unsigned int i, unsigned int j, unsigned int k, unsigned int l) const {
        return i * shape[1] * shape[2] * shape[3] + 
               j * shape[2] * shape[3] + 
               k * shape[3] + 
               l;
    }

    // 这个加法需要支持“单向广播”。
    // 具体来说，`others` 可以具有与 `this` 不同的形状，形状不同的维度长度必须为 1。
    // `others` 长度为 1 但 `this` 长度不为 1 的维度将发生广播计算。
    // 例如，`this` 形状为 `[1, 2, 3, 4]`，`others` 形状为 `[1, 2, 1, 4]`，
    // 则 `this` 与 `others` 相加时，3 个形状为 `[1, 2, 1, 4]` 的子张量各自与 `others` 对应项相加。
    // Tensor4D &operator+=(Tensor4D const &others) {
    //     // TODO: 实现单向广播的加法
    //     return *this;
    // }


    // 实现单向广播的加法
    Tensor4D &operator+=(Tensor4D const &other) {
        // 检查形状兼容性
        for (int dim = 0; dim < 4; ++dim) {
            if (shape[dim] != other.shape[dim] && other.shape[dim] != 1) {
                ASSERT(false, "Incompatible shapes for broadcasting");
                return *this;
            }
        }

        // 遍历当前张量的所有元素
        for (unsigned int i = 0; i < shape[0]; ++i) {
            for (unsigned int j = 0; j < shape[1]; ++j) {
                for (unsigned int k = 0; k < shape[2]; ++k) {
                    for (unsigned int l = 0; l < shape[3]; ++l) {
                        // 计算other张量的对应索引（考虑广播）
                        unsigned int other_i = (other.shape[0] == 1) ? 0 : i;
                        unsigned int other_j = (other.shape[1] == 1) ? 0 : j;
                        unsigned int other_k = (other.shape[2] == 1) ? 0 : k;
                        unsigned int other_l = (other.shape[3] == 1) ? 0 : l;
                        
                        unsigned int idx_this = linear_index(i, j, k, l);
                        unsigned int idx_other = other.linear_index(other_i, other_j, other_k, other_l);
                        
                        data[idx_this] += other.data[idx_other];
                    }
                }
            }
        }
        
        return *this;
    }
};

// ---- 不要修改以下代码 ----
int main(int argc, char **argv) {
    {
        unsigned int shape[]{1, 2, 3, 4};
        // clang-format off
        int data[]{
             1,  2,  3,  4,
             5,  6,  7,  8,
             9, 10, 11, 12,

            13, 14, 15, 16,
            17, 18, 19, 20,
            21, 22, 23, 24};
        // clang-format on
        auto t0 = Tensor4D(shape, data);
        auto t1 = Tensor4D(shape, data);
        t0 += t1;
        for (auto i = 0u; i < sizeof(data) / sizeof(*data); ++i) {
            ASSERT(t0.data[i] == data[i] * 2, "Tensor doubled by plus its self.");
        }
    }
    {
        unsigned int s0[]{1, 2, 3, 4};
        // clang-format off
        float d0[]{
            1, 1, 1, 1,
            2, 2, 2, 2,
            3, 3, 3, 3,

            4, 4, 4, 4,
            5, 5, 5, 5,
            6, 6, 6, 6};
        // clang-format on
        unsigned int s1[]{1, 2, 3, 1};
        // clang-format off
        float d1[]{
            6,
            5,
            4,

            3,
            2,
            1};
        // clang-format on

        auto t0 = Tensor4D(s0, d0);
        auto t1 = Tensor4D(s1, d1);
        t0 += t1;
        for (auto i = 0u; i < sizeof(d0) / sizeof(*d0); ++i) {
            ASSERT(t0.data[i] == 7.f, "Every element of t0 should be 7 after adding t1 to it.");
        }
    }
    {
        unsigned int s0[]{1, 2, 3, 4};
        // clang-format off
        double d0[]{
             1,  2,  3,  4,
             5,  6,  7,  8,
             9, 10, 11, 12,

            13, 14, 15, 16,
            17, 18, 19, 20,
            21, 22, 23, 24};
        // clang-format on
        unsigned int s1[]{1, 1, 1, 1};
        double d1[]{1};

        auto t0 = Tensor4D(s0, d0);
        auto t1 = Tensor4D(s1, d1);
        t0 += t1;
        for (auto i = 0u; i < sizeof(d0) / sizeof(*d0); ++i) {
            ASSERT(t0.data[i] == d0[i] + 1, "Every element of t0 should be incremented by 1 after adding t1 to it.");
        }
    }
}


/*
这么长

类模板
从函数模板升级了

四维，每一维的维度；指向四维数组的指针；
牛逼，666考广播机制




ai：

代码解释
1. 构造函数
Tensor4D(unsigned int const shape_[4], T const *data_) {
    unsigned int size = 1;
    for (int i = 0; i < 4; ++i) {
        shape[i] = shape_[i];
        size *= shape[i];
    }
    data = new T[size];
    std::memcpy(data, data_, size * sizeof(T));
}
复制shape数组
计算总元素数量（各维度乘积）
分配内存并复制数据
2. 线性索引计算
unsigned int linear_index(unsigned int i, unsigned int j, unsigned int k, unsigned int l) const {
    return i * shape[1] * shape[2] * shape[3] + 
           j * shape[2] * shape[3] + 
           k * shape[3] + 
           l;
}
将4D坐标(i, j, k, l)转换为线性数组索引。

3. 广播加法
Tensor4D &operator+=(Tensor4D const &other) {
    // 形状兼容性检查
    for (int dim = 0; dim < 4; ++dim) {
        if (shape[dim] != other.shape[dim] && other.shape[dim] != 1) {
            ASSERT(false, "Incompatible shapes for broadcasting");
            return *this;
        }
    }

    // 遍历所有元素
    for (unsigned int i = 0; i < shape[0]; ++i) {
        for (unsigned int j = 0; j < shape[1]; ++j) {
            for (unsigned int k = 0; k < shape[2]; ++k) {
                for (unsigned int l = 0; l < shape[3]; ++l) {
                    // 计算other的对应索引（考虑广播）
                    unsigned int other_i = (other.shape[0] == 1) ? 0 : i;
                    unsigned int other_j = (other.shape[1] == 1) ? 0 : j;
                    unsigned int other_k = (other.shape[2] == 1) ? 0 : k;
                    unsigned int other_l = (other.shape[3] == 1) ? 0 : l;
                    
                    unsigned int idx_this = linear_index(i, j, k, l);
                    unsigned int idx_other = other.linear_index(other_i, other_j, other_k, other_l);
                    
                    data[idx_this] += other.data[idx_other];
                }
            }
        }
    }
    
    return *this;
}
广播规则：

如果other的某个维度长度为1，则在该维度上进行广播
如果维度长度不相等且都不为1，则形状不兼容
广播示例
相同形状：[1,2,3,4] += [1,2,3,4] - 直接对应相加
部分广播：[1,2,3,4] += [1,2,3,1] - 第4维广播
完全广播：[1,2,3,4] += [1,1,1,1] - 所有维度广播
测试验证
程序运行后会通过所有测试：

相同形状张量相加正确
部分广播加法正确
完全广播加法正确
各种数据类型（int, float, double）都支持
这个实现提供了完整的4D张量支持，包括形状兼容性检查和广播机制，类似于NumPy和PyTorch中的广播功能。
*/